# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import numpy as np
from collections import namedtuple

from experimental.prophet.graph import graph_def_pb2


def to_data_type_proto(t):
  if t == float:
    return graph_def_pb2.DATA_TYPE_DOUBLE
  elif t == int:
    return graph_def_pb2.DATA_TYPE_INT64
  elif t == bool:
    return graph_def_pb2.DATA_TYPE_BOOL
  elif t == str:
    return graph_def_pb2.DATA_TYPE_STRING
  elif isinstance(t, int) and t in graph_def_pb2.DataType.values():
    return t
  else:
    raise ValueError('Unknown type: %s' % str(t))


def from_data_type_proto(dt):
  if dt == graph_def_pb2.DATA_TYPE_DOUBLE:
    return np.dtype(float)
  elif dt == graph_def_pb2.DATA_TYPE_INT64:
    return np.dtype(int)
  elif dt == graph_def_pb2.DATA_TYPE_BOOL:
    return np.dtype(bool)
  elif dt == graph_def_pb2.DATA_TYPE_STRING:
    return np.dtype(str)
  elif isinstance(dt, int) and dt >= graph_def_pb2.DATA_TYPE_OBJECT:
    return dt
  else:
    raise ValueError('Unknown dtype: %s' % str(dt))


_namedtuple_dict = {}


def _get_namedtuple_cls(fields):
  fields = tuple(fields)
  if fields in _namedtuple_dict:
    return _namedtuple_dict[fields]

  cls = namedtuple('Output', fields)
  _namedtuple_dict[fields] = cls
  return cls


def build_output_tuple(vars):
  output_tuple_cls = _get_namedtuple_cls([var.arg_name for var in vars])
  return output_tuple_cls(*vars)
